"""
WebSocket代理服务
================

负责WebSocket连接的代理转发和管理
支持负载均衡、故障转移、连接池管理
"""

import asyncio
import logging
import json
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
from datetime import datetime, timezone

import httpx
import websockets
from websockets.exceptions import WebSocketException
from fastapi import WebSocket, WebSocketDisconnect
from starlette.websockets import WebSocketState

from ..config import ServiceInstance, RouteConfig

logger = logging.getLogger(__name__)


@dataclass
class WebSocketRoute:
    """WebSocket路由配置"""
    path: str
    service_name: str
    target_path: str
    auth_required: bool = False
    priority: int = 10
    metadata: Dict[str, Any] = None


class WebSocketProxy:
    """WebSocket代理服务"""
    
    def __init__(self, services: Dict[str, List[ServiceInstance]], routes: List[WebSocketRoute]):
        self.services = services
        self.routes = {route.path: route for route in routes}
        self.active_connections: Dict[str, WebSocket] = {}
        self.connection_count = 0
        self.max_connections = 1000
        
    async def handle_websocket(self, websocket: WebSocket, path: str) -> bool:
        """
        处理WebSocket连接
        
        Args:
            websocket: WebSocket连接对象
            path: 请求路径
            
        Returns:
            bool: 处理是否成功
        """
        try:
            # 检查连接数限制
            if self.connection_count >= self.max_connections:
                logger.warning(f"WebSocket连接数已达上限: {self.max_connections}")
                await websocket.close(code=1013, reason="Server overloaded")
                return False
            
            # 查找匹配的路由
            route = self._find_route(path)
            if not route:
                logger.warning(f"未找到WebSocket路由: {path}")
                await websocket.close(code=1008, reason="Route not found")
                return False
            
            # 检查服务是否可用
            service_instances = self.services.get(route.service_name, [])
            if not service_instances:
                logger.error(f"服务 {route.service_name} 不可用")
                await websocket.close(code=1011, reason="Service unavailable")
                return False
            
            # 选择服务实例（简单轮询）
            service_instance = self._select_service_instance(service_instances)
            if not service_instance:
                logger.error(f"无法选择服务实例: {route.service_name}")
                await websocket.close(code=1011, reason="No available service instance")
                return False
            
            # 建立到后端服务的WebSocket连接
            backend_ws_url = self._build_websocket_url(service_instance, route.target_path)
            success = await self._proxy_websocket(websocket, backend_ws_url, path)
            
            return success
            
        except Exception as e:
            logger.error(f"WebSocket处理失败: {e}", exc_info=True)
            try:
                await websocket.close(code=1011, reason="Internal server error")
            except:
                pass
            return False
    
    def _find_route(self, path: str) -> Optional[WebSocketRoute]:
        """查找匹配的WebSocket路由"""
        # 精确匹配
        if path in self.routes:
            return self.routes[path]
        
        # 通配符匹配
        for route_path, route in self.routes.items():
            if route_path.endswith("*"):
                prefix = route_path[:-1]
                if path.startswith(prefix):
                    return route
        
        return None
    
    def _select_service_instance(self, instances: List[ServiceInstance]) -> Optional[ServiceInstance]:
        """选择服务实例（简单轮询算法）"""
        if not instances:
            return None
        
        # 过滤健康的实例
        healthy_instances = [inst for inst in instances if inst.is_healthy]
        if not healthy_instances:
            return instances[0]  # 如果没有健康实例，返回第一个
        
        # 简单轮询
        index = self.connection_count % len(healthy_instances)
        return healthy_instances[index]
    
    def _build_websocket_url(self, service_instance: ServiceInstance, target_path: str) -> str:
        """构建WebSocket URL"""
        protocol = "wss" if service_instance.ssl_enabled else "ws"
        return f"{protocol}://{service_instance.host}:{service_instance.port}{target_path}"
    
    async def _proxy_websocket(self, client_ws: WebSocket, backend_url: str, path: str) -> bool:
        """
        真正的WebSocket代理连接
        使用websockets库实现双向消息转发
        """
        connection_id = f"{path}_{self.connection_count}"
        self.connection_count += 1
        backend_ws = None
        
        try:
            # 1. 接受客户端连接
            await client_ws.accept()
            self.active_connections[connection_id] = client_ws
            logger.info(f"✅ 客户端连接已接受: {connection_id}")
            
            # 2. 连接到后端服务
            logger.info(f"🔗 正在连接后端服务: {backend_url}")
            backend_ws = await websockets.connect(
                backend_url,
                ping_interval=30,
                ping_timeout=10,
                close_timeout=10
            )
            logger.info(f"✅ 后端连接已建立: {backend_url}")
            
            # 3. 创建双向转发任务
            async def forward_to_backend():
                """客户端 → 后端"""
                try:
                    while True:
                        message = await client_ws.receive_text()
                        await backend_ws.send(message)
                        logger.debug(f"📤 转发到后端: {message[:100]}...")
                except WebSocketDisconnect:
                    logger.info(f"客户端 {connection_id} 断开")
                except Exception as e:
                    logger.error(f"转发到后端失败: {e}")
                    raise
            
            async def forward_to_client():
                """后端 → 客户端"""
                try:
                    async for message in backend_ws:
                        await client_ws.send_text(message)
                        logger.debug(f"📥 转发到客户端: {message[:100]}...")
                except Exception as e:
                    logger.error(f"转发到客户端失败: {e}")
                    raise
            
            # 4. 并发执行双向转发
            logger.info(f"🔄 开始双向消息转发: {connection_id}")
            await asyncio.gather(
                forward_to_backend(),
                forward_to_client(),
                return_exceptions=True
            )
            
            return True
            
        except websockets.exceptions.WebSocketException as e:
            logger.error(f"❌ WebSocket错误 [{connection_id}]: {e}")
            return False
        except Exception as e:
            logger.error(f"❌ 代理失败 [{connection_id}]: {e}", exc_info=True)
            return False
        finally:
            # 5. 清理资源
            if backend_ws:
                try:
                    await backend_ws.close()
                    logger.info(f"🔌 后端连接已关闭: {connection_id}")
                except:
                    pass
            if connection_id in self.active_connections:
                del self.active_connections[connection_id]
            self.connection_count = max(0, self.connection_count - 1)
            logger.info(f"🧹 连接 {connection_id} 已清理")
    
    
    def get_connection_count(self) -> int:
        """获取当前连接数"""
        return len(self.active_connections)
    
    def get_connection_info(self) -> Dict[str, Any]:
        """获取连接信息"""
        return {
            "active_connections": len(self.active_connections),
            "total_connections": self.connection_count,
            "max_connections": self.max_connections,
            "routes": list(self.routes.keys()),
            "services": list(self.services.keys())
        }
    
    async def close_all_connections(self):
        """关闭所有连接"""
        for connection_id, websocket in self.active_connections.items():
            try:
                if websocket.client_state != WebSocketState.DISCONNECTED:
                    await websocket.close(code=1000, reason="Server shutdown")
            except Exception as e:
                logger.warning(f"关闭连接 {connection_id} 失败: {e}")
        
        self.active_connections.clear()
        self.connection_count = 0
        logger.info("所有WebSocket连接已关闭")


# 全局WebSocket代理实例
websocket_proxy: Optional[WebSocketProxy] = None


def get_websocket_proxy() -> Optional[WebSocketProxy]:
    """获取WebSocket代理实例"""
    return websocket_proxy


def initialize_websocket_proxy(services: Dict[str, List[ServiceInstance]], routes: List[WebSocketRoute]):
    """初始化WebSocket代理"""
    global websocket_proxy
    websocket_proxy = WebSocketProxy(services, routes)
    logger.info("WebSocket代理已初始化")
